{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fc4bd651",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from chamferdist import ChamferDistance\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ae260630",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "source_cloud = torch.randn(1, 100, 3).cuda(2)\n",
    "target_cloud = torch.randn(1, 50, 3).cuda(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "89a5f15c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# source_cloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d6678e52",
   "metadata": {},
   "outputs": [],
   "source": [
    "chamferDist = ChamferDistance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7c5fae95",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "47.30221939086914\n"
     ]
    }
   ],
   "source": [
    "dist_forward = chamferDist(source_cloud, target_cloud)\n",
    "print(dist_forward.detach().cpu().item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "77f47512",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "17.50723648071289\n"
     ]
    }
   ],
   "source": [
    "dist_backward = chamferDist(source_cloud, target_cloud, reverse=True)\n",
    "print(dist_backward.detach().cpu().item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "30fcb1bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64.80945587158203\n"
     ]
    }
   ],
   "source": [
    "dist_bidirectional = chamferDist(source_cloud, target_cloud, bidirectional=True)\n",
    "print(dist_bidirectional.detach().cpu().item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84bebdc5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aff86e3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ed402e9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward Chamfer distance: 57.296119689941406\n",
      "Backward Chamfer distance: 16.604726791381836\n",
      "Backward Chamfer distance: 16.604726791381836\n",
      "Bi-directional Chamfer distance: 36.95042419433594\n",
      "Chamfer distance (self): 0.0\n",
      "Chamfer distance (self): 0.0\n",
      "Gradient norm wrt bidirectional Chamfer distance: 9.445826530456543\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import torch\n",
    "\n",
    "from chamferdist import ChamferDistance\n",
    "\n",
    "\n",
    "# Create two random pointclouds\n",
    "# (Batchsize x Number of points x Number of dims)\n",
    "source_cloud = torch.randn(1, 100, 3).cuda(2)\n",
    "target_cloud = torch.randn(1, 50, 3).cuda(2)\n",
    "source_cloud.requires_grad = True\n",
    "\n",
    "# Initialize Chamfer distance module\n",
    "chamferDist = ChamferDistance()\n",
    "# Compute Chamfer distance\n",
    "dist_forward = chamferDist(source_cloud, target_cloud)\n",
    "print(\"Forward Chamfer distance:\", dist_forward.detach().cpu().item())\n",
    "\n",
    "# Chamfer distance depends on the direction in which it is computed (as the\n",
    "# nearest neighbour varies, in each direction). One can either flip the order\n",
    "# of the arguments, or simply use the \"reverse\" flag.\n",
    "dist_backward = chamferDist(source_cloud, target_cloud, reverse=True)\n",
    "print(\"Backward Chamfer distance:\", dist_backward.detach().cpu().item())\n",
    "# Or, if you rather prefer, flip the order of the arguments.\n",
    "dist_backward = chamferDist(target_cloud, source_cloud)\n",
    "print(\"Backward Chamfer distance:\", dist_backward.detach().cpu().item())\n",
    "\n",
    "# To get a symmetric measure, the simplest way is to average both the \"forward\"\n",
    "# and \"backward\" distances. This is done by the \"bidirectional\" flag.\n",
    "cdist = 0.5 * chamferDist(source_cloud, target_cloud, bidirectional=True)\n",
    "cdist = 0.5 * chamferDist(target_cloud, source_cloud, bidirectional=True)\n",
    "print(\"Bi-directional Chamfer distance:\", cdist.detach().cpu().item())\n",
    "\n",
    "# As a sanity check, chamfer distance between a pointcloud and itself must be\n",
    "# zero.\n",
    "dist_self = chamferDist(source_cloud, source_cloud)\n",
    "print(\"Chamfer distance (self):\", dist_self.detach().cpu().item())\n",
    "dist_self = chamferDist(target_cloud, target_cloud)\n",
    "print(\"Chamfer distance (self):\", dist_self.detach().cpu().item())\n",
    "\n",
    "# Backprop using this loss!\n",
    "cdist.backward()\n",
    "print(\n",
    "    \"Gradient norm wrt bidirectional Chamfer distance:\",\n",
    "    source_cloud.grad.norm().detach().cpu().item(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40798d16",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6376ade",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
